"""
数据流管道。
"""

import numpy as np
import pandas as pd
import cv2
import os
import json
import paddle
from paddle.io import Dataset, DataLoader
from transform import random_transform
from utils import get_objectness_label


class vehicleDataset(Dataset):
    """
    车辆数据集。
    """

    def __init__(self, data_dict):
        """
        Args:
        ---
            data_dict: 列表格式，每个元素包含图片`img`与框`labels`。
                `labels`为装载[`xmin`, `ymin`, `xmax`, `ymax`]的列表。
        """
        super(vehicleDataset, self).__init__()
        self.data_dict = data_dict
        cfg = json.load(open('./config/cfg.json', 'r'))
        self.anchors = cfg['anchors']

    def __getitem__(self, idx):
        img, labels = random_transform(self.data_dict[idx]['img'], labels=self.data_dict[idx]['labels'])
        labels = get_objectness_label(np.array([img]), label_rects=np.array([labels]), anchors=self.anchors)
        # assert labels[0, :, 4, :, :].max() == 0
        return img.astype('float32'), labels[0]

    def __len__(self):
        return len(self.data_dict)


def build_dataset(mode='train'):
    rate = 0.8
    if mode == 'train' or mode == 'val':
        data_path = './data/training_images/'
        if mode == 'val':
            rate = 0.2
    else:
        data_path = './data/testing_images/'
        rate = 1
    label_path = './data/train_solution_bounding_boxes.csv'
    raw_label_data = pd.read_csv(label_path)
    img_data_files = os.listdir(data_path)
    print(f'reading from {data_path}')
    img_data = {}
    size = int(rate * len(img_data_files))
    if mode == 'train':
        img_data_files = img_data_files[:size]
    elif mode == 'val':
        img_data_files = img_data_files[-size:]
    for img_file in img_data_files:
        img = cv2.imread(os.path.join(data_path, img_file))
        img_data[img_file] = {
            'img': img,
            'labels': []
        }
    for i in raw_label_data.index:
        if raw_label_data.loc[i, 'image'] not in img_data:
            continue
        points = raw_label_data.iloc[i, 1:].values.tolist()
        points = [points[1], points[0], points[3], points[2]]
        img_data[raw_label_data.loc[i, 'image']]['labels'].append(points)
    vehicle_dataset = vehicleDataset(list(img_data.values()))
    return vehicle_dataset


if __name__ == '__main__':
    vehicle_dataset = build_dataset(mode='train')
